package me.tuyou.function.iterate;

import org.apache.flink.api.common.functions.FlatJoinFunction;
import org.apache.flink.api.common.functions.GroupReduceFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.*;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.util.Collector;

import java.util.Iterator;

/**
 * 批处理的增量迭代
 */
public class DeltaIterationTest {
    public static void main(String[] args) throws Exception {
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSource<Tuple2<Long, Double>> initSolutionSet = env.fromElements(new Tuple2<>(1L, 4.17d));
        DataSource<Tuple2<Long, Double>> initDeltaSet = env.fromElements(
                new Tuple2<>(1L, 1.9d),
                new Tuple2<>(2L, 4.8d),
                new Tuple2<>(3L, 2.9d)
        );

        DeltaIteration<Tuple2<Long, Double>, Tuple2<Long, Double>> iteration =
                initSolutionSet.iterateDelta(initDeltaSet, 100, 0);

        GroupReduceOperator<Tuple2<Long, Double>, Tuple2<Long, Double>> candidate = iteration
                .getWorkset()
                .groupBy(0)
                .reduceGroup(new GroupReduceFunction<Tuple2<Long, Double>, Tuple2<Long, Double>>() {
                    @Override
                    public void reduce(Iterable<Tuple2<Long, Double>> values, Collector<Tuple2<Long, Double>> out) throws Exception {
                        Iterator<Tuple2<Long, Double>> iter = values.iterator();
                        while (iter.hasNext()) {
                            Tuple2<Long, Double> item = iter.next();
                            out.collect(new Tuple2<>(item.f0, item.f1 + 1));
                        }
                    }
                });

        DataSet<Tuple2<Long, Double>> delta = candidate
                .join(iteration.getSolutionSet())
                .where(0)
                .equalTo(0)
                .with(new FlatJoinFunction<Tuple2<Long, Double>, Tuple2<Long, Double>, Tuple2<Long, Double>>() {
                    @Override
                    public void join(Tuple2<Long, Double> first, Tuple2<Long, Double> second, Collector<Tuple2<Long, Double>> out) throws Exception {
                        if (second != null) {
                            out.collect(new Tuple2<>(first.f0, first.f1 + second.f1));
                        } else {
                            out.collect(first);
                        }
                    }
                });

        FilterOperator<Tuple2<Long, Double>> nextWorkSet = delta.filter(i -> i.f0 < 1);
        iteration.closeWith(delta, nextWorkSet).print();
    }
}
